package com.rem.flink.flink8State;

import cn.hutool.core.collection.ListUtil;
import com.rem.flink.flink2Source.ClickSource;
import com.rem.flink.flink2Source.Event;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;

import java.util.ArrayList;
import java.util.List;

/**
 * 算子状态 -- 列表状态（ListState）
 * 每10个数据输出一次
 *
 * @author Rem
 * @date 2022-10-25
 */

public class OperatorListStateTest {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        SingleOutputStreamOperator<Event> stream = env.addSource(new ClickSource())
                .assignTimestampsAndWatermarks(WatermarkStrategy.<Event>forMonotonousTimestamps()
                        .withTimestampAssigner((element, recordTimestamp) -> element.getTimestamp()));

        stream.printToErr("time");

        // 批量缓存输出
        stream.addSink(new BufferingSink(10));

        env.execute();
    }

    private static class BufferingSink implements SinkFunction<Event>, CheckpointedFunction {
        //定义输出量
        int threshold;

        public BufferingSink(int threshold) {
            this.threshold = threshold;
            bufferedElements = new ArrayList<>();
        }

        ListState<Event> checkpointedSatre;
        List<Event> bufferedElements;

        /**
         * 正常情况下执行输出
         * 将给定值写入接收器。每个记录都调用这个函数。
         * 在实现SinkFunction时，您必须覆盖此方法，这是一个 default方法仅支持向后兼容旧式方法。
         *
         * @param value   输入记录。
         * @param context 输入记录的附加上下文。
         */
        @Override
        public void invoke(Event value, Context context) {
            bufferedElements.add(value);
            if (bufferedElements.size() == threshold) {
                // 输出到外部系统，这里用控制台打印模拟
                bufferedElements.forEach(System.out::println);
                System.out.println("==========输出完毕=========");
                bufferedElements.clear();
            }
        }

        /**
         * 将局部变量加入到检查点中
         * <p>
         * 当请求检查点的快照时，将调用此方法。这就像一个钩子 用于确保所有状态的函数是在初始化时通过FunctionInitializationContext提供的，
         * 或者现在通过FunctionSnapshotContext本身提供的。
         *
         * @param context
         * @throws Exception
         */
        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            checkpointedSatre.clear();
            checkpointedSatre.addAll(bufferedElements);
        }

        /**
         * 如果是从故障中恢复，就将ListState中的所有元素添加到局部变量中
         * <p>
         * 在分布式期间创建并行函数实例时调用此方法 执行。函数通常在此方法中设置其状态存储数据结构。
         *
         * @param context
         * @throws Exception
         */
        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            checkpointedSatre = context.getOperatorStateStore().getListState(new ListStateDescriptor<>("buffered-elements", Types.POJO(Event.class)));

            if (context.isRestored()) {
                bufferedElements = ListUtil.toList(checkpointedSatre.get());
            }
        }
    }
}
